{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9275270b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "词表大小: 10000\n"
     ]
    }
   ],
   "source": [
    "from tokenizers import Tokenizer\n",
    "\n",
    "# 加载 BPE Tokenizer\n",
    "tokenizer = Tokenizer.from_file(\"bpe_tokenizer.json\")\n",
    "\n",
    "# 获取词表大小\n",
    "vocab_size = tokenizer.get_vocab_size()\n",
    "print(\"词表大小:\", vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a55edba3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练数据大小: torch.Size([2308685, 50]) torch.Size([2308685, 50])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "def create_training_data(text, tokenizer, seq_length=50):\n",
    "    \"\"\"将文本转换为 Token ID，并生成训练样本\"\"\"\n",
    "    token_ids = tokenizer.encode(text).ids\n",
    "    inputs, targets = [], []\n",
    "\n",
    "    for i in range(len(token_ids) - seq_length):\n",
    "        inputs.append(token_ids[i:i+seq_length])\n",
    "        targets.append(token_ids[i+1:i+seq_length+1])\n",
    "\n",
    "    return torch.tensor(inputs), torch.tensor(targets)\n",
    "\n",
    "# 读取训练数据\n",
    "file_paths = [\"./processed_data.txt\"]\n",
    "text = \"\\n\".join([open(f, \"r\", encoding=\"utf-8\").read() for f in file_paths])\n",
    "\n",
    "# 生成训练样本\n",
    "inputs, targets = create_training_data(text, tokenizer)\n",
    "print(\"训练数据大小:\", inputs.shape, targets.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c12b117e",
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[4], line 108\u001b[0m\n\u001b[0;32m    105\u001b[0m     plt\u001b[38;5;241m.\u001b[39mshow()\n\u001b[0;32m    107\u001b[0m \u001b[38;5;66;03m# 训练模型\u001b[39;00m\n\u001b[1;32m--> 108\u001b[0m train(model, dataloader, epochs\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m20\u001b[39m)\n",
      "Cell \u001b[1;32mIn[4], line 81\u001b[0m, in \u001b[0;36mtrain\u001b[1;34m(model, dataloader, epochs)\u001b[0m\n\u001b[0;32m     79\u001b[0m     loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[0;32m     80\u001b[0m     torch\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mutils\u001b[38;5;241m.\u001b[39mclip_grad_norm_(model\u001b[38;5;241m.\u001b[39mparameters(), max_norm\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1.0\u001b[39m)  \u001b[38;5;66;03m# 防止梯度爆炸\u001b[39;00m\n\u001b[1;32m---> 81\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[0;32m     83\u001b[0m     epoch_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mitem()\n\u001b[0;32m     85\u001b[0m scheduler\u001b[38;5;241m.\u001b[39mstep()\n",
      "File \u001b[1;32md:\\anaconda3\\Lib\\site-packages\\torch\\optim\\lr_scheduler.py:137\u001b[0m, in \u001b[0;36mLRScheduler.__init__.<locals>.patch_track_step_called.<locals>.wrap_step.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    135\u001b[0m opt \u001b[38;5;241m=\u001b[39m opt_ref()\n\u001b[0;32m    136\u001b[0m opt\u001b[38;5;241m.\u001b[39m_opt_called \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m  \u001b[38;5;66;03m# type: ignore[union-attr]\u001b[39;00m\n\u001b[1;32m--> 137\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m func\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__get__\u001b[39m(opt, opt\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m)(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32md:\\anaconda3\\Lib\\site-packages\\torch\\optim\\optimizer.py:487\u001b[0m, in \u001b[0;36mOptimizer.profile_hook_step.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    482\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    483\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[0;32m    484\u001b[0m                 \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfunc\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m must return None or a tuple of (new_args, new_kwargs), but got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mresult\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m    485\u001b[0m             )\n\u001b[1;32m--> 487\u001b[0m out \u001b[38;5;241m=\u001b[39m func(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m    488\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_optimizer_step_code()\n\u001b[0;32m    490\u001b[0m \u001b[38;5;66;03m# call optimizer step post hooks\u001b[39;00m\n",
      "File \u001b[1;32md:\\anaconda3\\Lib\\site-packages\\torch\\optim\\optimizer.py:91\u001b[0m, in \u001b[0;36m_use_grad_for_differentiable.<locals>._use_grad\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m     89\u001b[0m     torch\u001b[38;5;241m.\u001b[39mset_grad_enabled(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdefaults[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdifferentiable\u001b[39m\u001b[38;5;124m\"\u001b[39m])\n\u001b[0;32m     90\u001b[0m     torch\u001b[38;5;241m.\u001b[39m_dynamo\u001b[38;5;241m.\u001b[39mgraph_break()\n\u001b[1;32m---> 91\u001b[0m     ret \u001b[38;5;241m=\u001b[39m func(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m     92\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m     93\u001b[0m     torch\u001b[38;5;241m.\u001b[39m_dynamo\u001b[38;5;241m.\u001b[39mgraph_break()\n",
      "File \u001b[1;32md:\\anaconda3\\Lib\\site-packages\\torch\\optim\\adam.py:223\u001b[0m, in \u001b[0;36mAdam.step\u001b[1;34m(self, closure)\u001b[0m\n\u001b[0;32m    211\u001b[0m     beta1, beta2 \u001b[38;5;241m=\u001b[39m group[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbetas\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[0;32m    213\u001b[0m     has_complex \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_init_group(\n\u001b[0;32m    214\u001b[0m         group,\n\u001b[0;32m    215\u001b[0m         params_with_grad,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    220\u001b[0m         state_steps,\n\u001b[0;32m    221\u001b[0m     )\n\u001b[1;32m--> 223\u001b[0m     adam(\n\u001b[0;32m    224\u001b[0m         params_with_grad,\n\u001b[0;32m    225\u001b[0m         grads,\n\u001b[0;32m    226\u001b[0m         exp_avgs,\n\u001b[0;32m    227\u001b[0m         exp_avg_sqs,\n\u001b[0;32m    228\u001b[0m         max_exp_avg_sqs,\n\u001b[0;32m    229\u001b[0m         state_steps,\n\u001b[0;32m    230\u001b[0m         amsgrad\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mamsgrad\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    231\u001b[0m         has_complex\u001b[38;5;241m=\u001b[39mhas_complex,\n\u001b[0;32m    232\u001b[0m         beta1\u001b[38;5;241m=\u001b[39mbeta1,\n\u001b[0;32m    233\u001b[0m         beta2\u001b[38;5;241m=\u001b[39mbeta2,\n\u001b[0;32m    234\u001b[0m         lr\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlr\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    235\u001b[0m         weight_decay\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mweight_decay\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    236\u001b[0m         eps\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124meps\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    237\u001b[0m         maximize\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmaximize\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    238\u001b[0m         foreach\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mforeach\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    239\u001b[0m         capturable\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcapturable\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    240\u001b[0m         differentiable\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdifferentiable\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    241\u001b[0m         fused\u001b[38;5;241m=\u001b[39mgroup[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfused\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m    242\u001b[0m         grad_scale\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgrad_scale\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m),\n\u001b[0;32m    243\u001b[0m         found_inf\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfound_inf\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m),\n\u001b[0;32m    244\u001b[0m     )\n\u001b[0;32m    246\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m loss\n",
      "File \u001b[1;32md:\\anaconda3\\Lib\\site-packages\\torch\\optim\\optimizer.py:154\u001b[0m, in \u001b[0;36m_disable_dynamo_if_unsupported.<locals>.wrapper.<locals>.maybe_fallback\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    152\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m disabled_func(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m    153\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 154\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m func(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32md:\\anaconda3\\Lib\\site-packages\\torch\\optim\\adam.py:784\u001b[0m, in \u001b[0;36madam\u001b[1;34m(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, foreach, capturable, differentiable, fused, grad_scale, found_inf, has_complex, amsgrad, beta1, beta2, lr, weight_decay, eps, maximize)\u001b[0m\n\u001b[0;32m    781\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    782\u001b[0m     func \u001b[38;5;241m=\u001b[39m _single_tensor_adam\n\u001b[1;32m--> 784\u001b[0m func(\n\u001b[0;32m    785\u001b[0m     params,\n\u001b[0;32m    786\u001b[0m     grads,\n\u001b[0;32m    787\u001b[0m     exp_avgs,\n\u001b[0;32m    788\u001b[0m     exp_avg_sqs,\n\u001b[0;32m    789\u001b[0m     max_exp_avg_sqs,\n\u001b[0;32m    790\u001b[0m     state_steps,\n\u001b[0;32m    791\u001b[0m     amsgrad\u001b[38;5;241m=\u001b[39mamsgrad,\n\u001b[0;32m    792\u001b[0m     has_complex\u001b[38;5;241m=\u001b[39mhas_complex,\n\u001b[0;32m    793\u001b[0m     beta1\u001b[38;5;241m=\u001b[39mbeta1,\n\u001b[0;32m    794\u001b[0m     beta2\u001b[38;5;241m=\u001b[39mbeta2,\n\u001b[0;32m    795\u001b[0m     lr\u001b[38;5;241m=\u001b[39mlr,\n\u001b[0;32m    796\u001b[0m     weight_decay\u001b[38;5;241m=\u001b[39mweight_decay,\n\u001b[0;32m    797\u001b[0m     eps\u001b[38;5;241m=\u001b[39meps,\n\u001b[0;32m    798\u001b[0m     maximize\u001b[38;5;241m=\u001b[39mmaximize,\n\u001b[0;32m    799\u001b[0m     capturable\u001b[38;5;241m=\u001b[39mcapturable,\n\u001b[0;32m    800\u001b[0m     differentiable\u001b[38;5;241m=\u001b[39mdifferentiable,\n\u001b[0;32m    801\u001b[0m     grad_scale\u001b[38;5;241m=\u001b[39mgrad_scale,\n\u001b[0;32m    802\u001b[0m     found_inf\u001b[38;5;241m=\u001b[39mfound_inf,\n\u001b[0;32m    803\u001b[0m )\n",
      "File \u001b[1;32md:\\anaconda3\\Lib\\site-packages\\torch\\optim\\adam.py:430\u001b[0m, in \u001b[0;36m_single_tensor_adam\u001b[1;34m(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, grad_scale, found_inf, amsgrad, has_complex, beta1, beta2, lr, weight_decay, eps, maximize, capturable, differentiable)\u001b[0m\n\u001b[0;32m    428\u001b[0m         denom \u001b[38;5;241m=\u001b[39m (max_exp_avg_sqs[i]\u001b[38;5;241m.\u001b[39msqrt() \u001b[38;5;241m/\u001b[39m bias_correction2_sqrt)\u001b[38;5;241m.\u001b[39madd_(eps)\n\u001b[0;32m    429\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 430\u001b[0m         denom \u001b[38;5;241m=\u001b[39m (exp_avg_sq\u001b[38;5;241m.\u001b[39msqrt() \u001b[38;5;241m/\u001b[39m bias_correction2_sqrt)\u001b[38;5;241m.\u001b[39madd_(eps)\n\u001b[0;32m    432\u001b[0m     param\u001b[38;5;241m.\u001b[39maddcdiv_(exp_avg, denom, value\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39mstep_size)\n\u001b[0;32m    434\u001b[0m \u001b[38;5;66;03m# Lastly, switch back to complex view\u001b[39;00m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m在当前单元格或上一个单元格中执行代码时 Kernel 崩溃。\n",
      "\u001b[1;31m请查看单元格中的代码，以确定故障的可能原因。\n",
      "\u001b[1;31m单击<a href='https://aka.ms/vscodeJupyterKernelCrash'>此处</a>了解详细信息。\n",
      "\u001b[1;31m有关更多详细信息，请查看 Jupyter <a href='command:jupyter.viewOutput'>log</a>。"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# 自动选择设备\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# 创建 PyTorch 数据加载器\n",
    "dataset = TensorDataset(inputs, targets)\n",
    "dataloader = DataLoader(dataset, batch_size=16, shuffle=True)\n",
    "\n",
    "# 定义 Transformer Decoder\n",
    "class TransformerDecoder(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, num_heads, num_layers, dim_feedforward, max_len):\n",
    "        \"\"\"\n",
    "        Transformer Decoder 网络架构\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model)\n",
    "        self.positional_encoding = nn.Parameter(torch.randn(1, max_len, d_model))\n",
    "\n",
    "        decoder_layer = nn.TransformerDecoderLayer(d_model, num_heads, dim_feedforward)\n",
    "        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers)\n",
    "        self.fc_out = nn.Linear(d_model, vocab_size)\n",
    "\n",
    "    def forward(self, input_seq, memory):\n",
    "        \"\"\"\n",
    "        input_seq: (batch_size, seq_len)\n",
    "        memory: (batch_size, seq_len)\n",
    "        return: (batch_size, seq_len, vocab_size)\n",
    "        \"\"\"\n",
    "        seq_len = input_seq.size(1)\n",
    "        embedded = self.embedding(input_seq) + self.positional_encoding[:, :seq_len, :]\n",
    "        memory_emb = self.embedding(memory) + self.positional_encoding[:, :memory.size(1), :]\n",
    "\n",
    "        # 注意: Transformer 模块需要 (seq_len, batch_size, d_model)\n",
    "        embedded = embedded.transpose(0, 1)\n",
    "        memory_emb = memory_emb.transpose(0, 1)\n",
    "\n",
    "        output = self.transformer_decoder(embedded, memory_emb)  # (seq_len, batch_size, d_model)\n",
    "        output = output.transpose(0, 1)  # (batch_size, seq_len, d_model)\n",
    "        return self.fc_out(output)  # (batch_size, seq_len, vocab_size)\n",
    "\n",
    "# 超参数\n",
    "d_model = 64\n",
    "num_heads = 4\n",
    "num_layers = 3\n",
    "dim_feedforward = 128\n",
    "max_len = 100\n",
    "\n",
    "model = TransformerDecoder(vocab_size, d_model, num_heads, num_layers, dim_feedforward, max_len).to(device)\n",
    "\n",
    "# 训练参数\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.95)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "def train(model, dataloader, epochs=20):\n",
    "    epochses = []\n",
    "    losses = []\n",
    "\n",
    "    plt.ion()  # 打开交互式绘图\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        epoch_loss = 0\n",
    "        for batch_inputs, batch_targets in dataloader:\n",
    "            batch_inputs = batch_inputs.to(device)\n",
    "            batch_targets = batch_targets.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            output = model(batch_inputs, batch_inputs)  # decoder 用自己做 memory\n",
    "\n",
    "            # output: (batch_size, seq_len, vocab_size)\n",
    "            loss = loss_fn(output.view(-1, vocab_size), batch_targets.view(-1))\n",
    "\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 防止梯度爆炸\n",
    "            optimizer.step()\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "\n",
    "        scheduler.step()\n",
    "        avg_loss = epoch_loss / len(dataloader)\n",
    "        print(f\"Epoch {epoch + 1}/{epochs}, Loss: {avg_loss:.4f}\")\n",
    "\n",
    "        epochses.append(epoch + 1)\n",
    "        losses.append(avg_loss)\n",
    "\n",
    "        # 实时更新可视化\n",
    "        ax.clear()\n",
    "        ax.plot(epochses, losses, label='Loss')\n",
    "        ax.set_xlabel('Epochs')\n",
    "        ax.set_ylabel('Loss')\n",
    "        ax.set_title('Training Loss')\n",
    "        ax.legend()\n",
    "        plt.pause(0.1)\n",
    "\n",
    "        # 保存模型\n",
    "        torch.save(model.state_dict(), f\"model_epoch_{epoch + 1}.pth\")\n",
    "\n",
    "    plt.ioff()\n",
    "    plt.show()\n",
    "\n",
    "# 训练模型\n",
    "train(model, dataloader, epochs=20)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a43b3943",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_text(model, tokenizer, start_text, max_length=50):\n",
    "    \"\"\"用 BPE 词表生成文本\"\"\"\n",
    "\n",
    "    # 加载训练好的模型\n",
    "    model.load_state_dict(torch.load(\"./model_epoch_49.pth\", weights_only=True))\n",
    "    \n",
    "    # 为什么要设置为评估模式？\n",
    "    # 因为在评估模式下，模型会关闭 Dropout 和 BatchNorm 等训练时使用的操作，从而使得模型在推理时更加稳定和一致。\n",
    "    model.eval()\n",
    "\n",
    "    input_tokens = tokenizer.encode(\"<bos> \" + start_text).ids\n",
    "    input_tensor = torch.tensor([input_tokens])\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_length):\n",
    "            output = model(input_tensor, input_tensor)\n",
    "            next_token = output.argmax(dim=-1)[:, -1].item()\n",
    "            input_tokens.append(next_token)\n",
    "            input_tensor = torch.tensor([input_tokens])\n",
    "\n",
    "            # 如果遇到 <eos> 结束标记，则停止\n",
    "            if tokenizer.decode([next_token]) == \"<eos>\":\n",
    "                break\n",
    "\n",
    "    return tokenizer.decode(input_tokens)\n",
    "\n",
    "# 测试生成\n",
    "text = \"hello ,how are you?\"\n",
    "\n",
    "word_count = len(text.split())  # 获取单词个数\n",
    "print(\"生成文本:\", generate_text(model, tokenizer, text, max_length=word_count))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "83f995b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "词表大小: 10000\n",
      "训练数据大小: torch.Size([658, 50]) torch.Size([658, 50])\n",
      "使用设备: cpu\n",
      "max_len: 50\n",
      "Epoch 1/50, Loss: 7.3441\n"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m在当前单元格或上一个单元格中执行代码时 Kernel 崩溃。\n",
      "\u001b[1;31m请查看单元格中的代码，以确定故障的可能原因。\n",
      "\u001b[1;31m单击<a href='https://aka.ms/vscodeJupyterKernelCrash'>此处</a>了解详细信息。\n",
      "\u001b[1;31m有关更多详细信息，请查看 Jupyter <a href='command:jupyter.viewOutput'>log</a>。"
     ]
    }
   ],
   "source": [
    "from tokenizers import Tokenizer\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# ----------- 数据部分 -----------\n",
    "# 加载 BPE Tokenizer\n",
    "tokenizer = Tokenizer.from_file(\"bpe_tokenizer.json\")\n",
    "\n",
    "# 获取词表大小\n",
    "vocab_size = tokenizer.get_vocab_size()\n",
    "print(\"词表大小:\", vocab_size)\n",
    "\n",
    "def create_training_data(text, tokenizer, seq_length=50):\n",
    "    \"\"\"将文本转换为 Token ID，并生成训练样本\"\"\"\n",
    "    token_ids = tokenizer.encode(text).ids\n",
    "    inputs, targets = [], []\n",
    "\n",
    "    for i in range(len(token_ids) - seq_length):\n",
    "        inputs.append(token_ids[i:i+seq_length])\n",
    "        targets.append(token_ids[i+1:i+seq_length+1])\n",
    "\n",
    "    return torch.tensor(inputs), torch.tensor(targets)\n",
    "\n",
    "# 读取训练数据\n",
    "file_paths = [\"data1.txt\", \"data2.txt\"]\n",
    "text = \"\\n\".join([open(f, \"r\", encoding=\"utf-8\").read() for f in file_paths])\n",
    "\n",
    "# 生成训练样本\n",
    "inputs, targets = create_training_data(text, tokenizer)\n",
    "print(\"训练数据大小:\", inputs.shape, targets.shape)\n",
    "\n",
    "# ----------- 训练数据 DataLoader -----------\n",
    "dataset = TensorDataset(inputs, targets)\n",
    "dataloader = DataLoader(dataset, batch_size=16, shuffle=True)\n",
    "\n",
    "# 自动设备检测\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"使用设备:\", device)\n",
    "\n",
    "# ----------- 模型定义 -----------\n",
    "class TransformerDecoder(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, num_heads, num_layers, dim_feedforward, max_len):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model)\n",
    "        self.positional_encoding = nn.Parameter(torch.randn(1, max_len, d_model))\n",
    "\n",
    "        decoder_layer = nn.TransformerDecoderLayer(d_model, num_heads, dim_feedforward)\n",
    "        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers)\n",
    "        self.fc_out = nn.Linear(d_model, vocab_size)\n",
    "\n",
    "    def forward(self, input_seq, memory):\n",
    "        seq_len = input_seq.size(1)\n",
    "        embedded = self.embedding(input_seq) + self.positional_encoding[:, :seq_len, :]\n",
    "        memory_emb = self.embedding(memory) + self.positional_encoding[:, :memory.size(1), :]\n",
    "\n",
    "        # Transformer 需要 (seq_len, batch_size, d_model)\n",
    "        embedded = embedded.transpose(0, 1)\n",
    "        memory_emb = memory_emb.transpose(0, 1)\n",
    "\n",
    "        output = self.transformer_decoder(embedded, memory_emb)\n",
    "        output = output.transpose(0, 1)  # (batch_size, seq_len, d_model)\n",
    "        return self.fc_out(output)\n",
    "\n",
    "# ----------- 超参数 -----------\n",
    "d_model = 64\n",
    "num_heads = 4\n",
    "num_layers = 3\n",
    "dim_feedforward = 128\n",
    "max_len = inputs.shape[1]  # 自动用 seq_length\n",
    "print(\"max_len:\", max_len)\n",
    "\n",
    "model = TransformerDecoder(vocab_size, d_model, num_heads, num_layers, dim_feedforward, max_len).to(device)\n",
    "\n",
    "# ----------- 优化器、调度器、损失函数 -----------\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.95)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "# ----------- 训练函数 -----------\n",
    "def train(model, dataloader, epochs=20):\n",
    "    epochses = []\n",
    "    losses = []\n",
    "\n",
    "    plt.ion()  # 打开交互式\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        epoch_loss = 0\n",
    "        for batch_inputs, batch_targets in dataloader:\n",
    "            batch_inputs = batch_inputs.to(device)\n",
    "            batch_targets = batch_targets.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            output = model(batch_inputs, batch_inputs)  # decoder memory = 输入\n",
    "\n",
    "            # output: (batch_size, seq_len, vocab_size)\n",
    "            # 需要 flatten 两个 tensor 对齐\n",
    "            loss = loss_fn(output.view(-1, vocab_size), batch_targets.view(-1))\n",
    "\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "            optimizer.step()\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "\n",
    "        scheduler.step()\n",
    "        avg_loss = epoch_loss / len(dataloader)\n",
    "        print(f\"Epoch {epoch + 1}/{epochs}, Loss: {avg_loss:.4f}\")\n",
    "\n",
    "        epochses.append(epoch + 1)\n",
    "        losses.append(avg_loss)\n",
    "\n",
    "        # 实时绘图\n",
    "        ax.clear()\n",
    "        ax.plot(epochses, losses, label='Loss')\n",
    "        ax.set_xlabel('Epochs')\n",
    "        ax.set_ylabel('Loss')\n",
    "        ax.set_title('Training Loss')\n",
    "        ax.legend()\n",
    "        plt.pause(0.1)\n",
    "\n",
    "        # 保存模型\n",
    "        torch.save(model.state_dict(), f\"model_epoch_{epoch + 1}.pth\")\n",
    "\n",
    "    plt.ioff()\n",
    "    plt.show()\n",
    "\n",
    "# ----------- 启动训练 -----------\n",
    "train(model, dataloader, epochs=50)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
